import queue
import threading
import time

from .priority_queue import Item, PriorityQueue


class DelayQueue:
    wakeup_event: threading.Event
    pq: PriorityQueue
    lock: threading
    timeout_queue: queue.Queue  # 超时的任务队列

    def now(self) -> int:
        return int(time.time() * 1000)

    def __init__(self, qsize: int) -> None:
        self.wakeup_event = threading.Event()
        self.pq = PriorityQueue(qsize)
        self.lock = threading.Lock()
        self.timeout_queue = queue.Queue()

    def wakup(self):
        self.wakeup_event.set()

    def offer(self, elem: any, expiration: int) -> bool:
        item = Item(value=elem, priority=expiration)
        with self.lock:
            ok = self.pq.insert(item)
            index = item.index
            if index == 0:
                # 有更早的超时任务添加了
                if not self.wakeup_event.is_set():
                    self.wakeup_event.set()
            self.wakeup_event.set()
            return ok

    def poll(self, exit_event: threading.Event):
        while not exit_event.is_set():
            current_timestamp = self.now()
            with self.lock:
                item = self.pq.peek_and_shif(current_timestamp)
                if item.item is None:
                    self.wakeup_event.clear()
            if item.item is None:
                if item.delta == 0:
                    # 队列没有元素, 等待插入元素唤醒
                    self.wakeup_event.wait()
                elif item.delta > 0:
                    self.wakeup_event.wait(timeout=item.delta)
            else:
                self.timeout_queue.put(item.item.value)

    def get(self) -> Item:
        return self.timeout_queue.get()


if __name__ == "__main__":
    dq = DelayQueue(1)
    now = int(time.time())
    now = 1666940000
    assert dq.offer("task1", now - 100) is True
    assert dq.offer("task2", now - 10) is False

    dq = DelayQueue(10)
    now = int(time.time())
    now = 1666940000
    dq.offer("task1", now - 100)
    dq.offer("task2", now - 10)
    dq.offer("task3", now + 2)
    dq.offer("task4", now + 1)
    dq.offer("task5", now + 3)
    dq.offer("task6", now + 4)
    dq.offer("task7", now + 0)
    item = dq.pq.peek_and_shif(now).item
    assert item.value == "task1", item.__dict__
    item = dq.pq.peek_and_shif(now).item
    assert item.value == "task2", item.__dict__
    item = dq.pq.peek_and_shif(now + 0).item
    assert item.value == "task7", (item, dq.timeout_queue.qsize(), item, dq.pq.get_data())
    item = dq.pq.peek_and_shif(now + 1).item
    assert item.value == "task4", item.__dict__
    item = dq.pq.peek_and_shif(now + 2).item
    assert item.value == "task3", (item, dq.timeout_queue.qsize(), item, dq.pq.get_data())
    item = dq.pq.peek_and_shif(now + 3).item
    assert item.value == "task5", (item, dq.timeout_queue.qsize(), item, dq.pq.get_data())
    item = dq.pq.peek_and_shif(now + 4).item
    assert item.value == "task6", (item, dq.timeout_queue.qsize(), item, dq.pq.get_data())
